#include "codec.h"



static std::string format (const char* fmt, ...) {
  c8 buf[4096];
  va_list args;
  va_start(args, fmt);
  vsnprintf(buf, 4096, fmt, args);
  va_end(args);
  return {buf};
}

static u16 checksum (const u8* data, uwl size) {
  u32 checksum = 0;
  auto wordp = (u16*)data;
  auto words = size >> 1;
  for (auto i = 0UL; i < words; ++i)
    checksum += wordp[i];
  if (size & 1)
    checksum += data[size - 1];
  checksum = (checksum & 0xFFFF) + (checksum >> 16);
  checksum += (checksum >> 16);
  return ~(u16)checksum;
}


static u16 checksum (const u16* head, const u8* data, uwl size) {
  u32 checksum = 0;
  for (auto i = 0; i < 6; ++i)
    checksum += head[i];
  auto wordp = (u16*)data;
  auto words = size >> 1;
  for (auto i = 0UL; i < words; ++i)
    checksum += wordp[i];
  if (size & 1)
    checksum += data[size - 1];
  checksum = (checksum & 0xFFFF) + (checksum >> 16);
  checksum += (checksum >> 16);
  return ~(u16)checksum;
}


void Codec::gen (u8 key[32]) {
  Random random;
  for (auto i = 0; i < 32; ++i)
    key[i] = random.uniform_int<uint8_t>(0, 255);
}


void Codec::nat (Ptr<IpPackage> pkg) {
  auto buf = pkg->buf;
  buf.writeIndex(10);
  buf.writeWord<u16>(0);
  buf.writeDWord(pkg->srcIPv4);
  buf.writeDWord(pkg->dstIPv4);
  buf.move(-10);
  buf.writeWord(htons(checksum(buf.base(), pkg->hdrSize)));
  buf.writeIndex(pkg->hdrSize);
  buf.writeWord(pkg->srcPort);
  buf.writeWord(pkg->dstPort);
  switch (pkg->protocol) {
    default:
    case IPPROTO_TCP: buf.writeIndex(pkg->hdrSize + 16); break;
    case IPPROTO_UDP: buf.writeIndex(pkg->hdrSize + 6); break;
  }
  buf.writeWord<u16>(0);
  buf.move(-2);
  auto payloadSize = pkg->pkgSize - pkg->hdrSize;
  struct Pseudo pseudo = {
    .srcIPv4 = htonl(pkg->srcIPv4),
    .dstIPv4 = htonl(pkg->dstIPv4),
    .reserved = 0,
    .protocol = pkg->protocol,
    .datSize = htons(payloadSize),
  };
  buf.writeWord<u16>(htons(checksum((u16*)&pseudo, buf.base() + pkg->hdrSize, payloadSize)));
}


Ptr<IpPackage> Codec::parse (ByteBuffer<>& buf) {
  while (buf.readable()) {
    auto versionAndHeaderLength = buf.readByte();
    auto version = versionAndHeaderLength >> 4;
    if (version != 4) {
      log_d("IP版本错误");
      continue;
    }
    auto hdrSize = (u16)((versionAndHeaderLength & 0xF) << 2);
    if (hdrSize < 20) {
      log_d("IP头长度错误");
      continue;
    }
    if (buf.readable() < hdrSize) {
      buf.skip(-1);
      return nullptr;
    }
    buf.skip(1);
    auto pkgSize = buf.readWord();
    if (pkgSize < hdrSize) {
      buf.skip(-1);
      log_d("IP包长度错误");
      continue;
    }
    if (0 != checksum(buf.readIter() - 4, hdrSize)) {
      buf.skip(-3);
      log_d("IP头校验错误");
      continue;
    }
    if (buf.readable() < (uwl)pkgSize - 4) {
      buf.skip(-4);
      return nullptr;
    }
    buf.skip(4);
    auto ttl = buf.readByte();
    auto protocol = buf.readByte();
    buf.skip(2);
    auto srcIPv4 = buf.readDWord();
    auto dstIPv4 = buf.readDWord();
    buf.skip(hdrSize - 20);
    auto payloadSize = pkgSize - hdrSize;
//    Pseudo pseudo = {
//      .srcIPv4 = htonl(srcIPv4),
//      .dstIPv4 = htonl(dstIPv4),
//      .reserved = 0,
//      .protocol = protocol,
//      .datSize = htons(payloadSize),
//    };
//    if (0 != checksum((u16*)&pseudo, buf.readIter(), payloadSize)) {
//      buf.skip(payloadSize);
//      pkg->err = "Checksum not matches";
//      return pkg;
//    }
    auto srcPort = buf.readWord();
    auto dstPort = buf.readWord();
    auto pkg = Ptr<IpPackage>::create();
    pkg->pkgSize = pkgSize;
    pkg->hdrSize = hdrSize;
    pkg->srcIPv4 = srcIPv4;
    pkg->dstIPv4 = dstIPv4;
    pkg->srcPort = srcPort;
    pkg->dstPort = dstPort;
    pkg->protocol = protocol;
    pkg->ttl = ttl;
    if (protocol == IPPROTO_TCP) {
      buf.skip(9);
      pkg->tcpFlag = buf.readByte();
      buf.skip(-10);
    }
    buf.skip(payloadSize - 4);
    pkg->buf = ByteBuffer<>::create(buf.readIter() - pkgSize, pkgSize);
    if (protocol != IPPROTO_TCP && protocol != IPPROTO_UDP)
      pkg->err = format("Unsupported protocol %d", protocol);
    if (!srcIPv4)
      pkg->err = "Invalid src address";
    if (!dstIPv4)
      pkg->err = "Invalid dst address";
    if (!srcPort)
      pkg->err = "Invalid src port";
    if (!dstPort)
      pkg->err = "Invalid dst port";
    return pkg;
  }
  return nullptr;
}


bool Codec::check (ByteBuffer<> buf, const char* prefix) {
  auto data = buf;
  auto dataSize = (u16)data.readable();
  if (dataSize < 20) {
    log_e("%s: 不足20字节: %d", prefix, dataSize);
    return false;
  }
  auto versionAndHeaderLength = data.readByte();
  auto version = versionAndHeaderLength >> 4;
  if (version != 4) {
    log_e("%s: 不支持的网络层协议: %d", prefix, version);
    return false;
  }
  auto hdrSize = (u16)((versionAndHeaderLength & 0xF) << 2);
  if (hdrSize < 20) {
    log_e("%s: IP头长度错误: %d", prefix, hdrSize);
    return false;
  }
  data.skip(1);
  auto pkgSize = data.readWord();
  if (pkgSize < hdrSize) {
    log_e("%s: IP包长度错误: %d", prefix, pkgSize);
    return false;
  }
  if (pkgSize > dataSize) {
    log_e("%s: IP包不完整: %d[%d]", prefix, pkgSize, dataSize);
    return false;
  }
  auto payloadSize = pkgSize - hdrSize;
  data.skip(5);
  auto protocol = data.readByte();
  if (protocol != IPPROTO_TCP && protocol != IPPROTO_UDP) {
    log_e("%s: 不支持的传输层协议: %d", prefix, protocol);
    return false;
  }
  data.skip(2);
  auto srcIPv4 = data.readDWord();
  auto dstIPv4 = data.readDWord();
  if (!srcIPv4) {
    log_e("%s: 无效的源地址: %s", prefix, Codec::ipaddr(srcIPv4).c_str());
    return false;
  }
  if (!dstIPv4) {
    log_e("%s: 无效的目的地址: %s", prefix, Codec::ipaddr(dstIPv4).c_str());
    return false;
  }
  data.skip(hdrSize - 20);
  auto srcPort = data.readWord();
  auto dstPort = data.readWord();
  if (!srcPort) {
    log_e("%s: 无效的源端口: %d", prefix, srcPort);
    return false;
  }
  if (!dstPort) {
    log_e("%s: 无效的目的端口: %d", prefix, dstPort);
    return false;
  }
  u16 cs = checksum(data.base(), hdrSize);
  if (cs) {
    log_e("%s: IP头校验错误: %x", prefix, cs);
    return false;
  }
  Pseudo pseudo = {
    .srcIPv4 = htonl(srcIPv4),
    .dstIPv4 = htonl(dstIPv4),
    .reserved = 0,
    .protocol = protocol,
    .datSize = htons(payloadSize),
  };
  cs = checksum((u16*)&pseudo, data.base() + hdrSize, payloadSize);
  if (cs) {
    log_e("%s: TCP/UDP校验错误: %x", prefix, cs);
    return false;
  }
  return true;
}


void Codec::print (Ptr<IpPackage> pkg, const char* prefix) {
  auto protocol = pkg->protocol == IPPROTO_TCP ? "TCP" : "UDP";
  std::string flags;
  if (pkg->tcpFlag & 0b00000001)
    flags += " FIN";
  if (pkg->tcpFlag & 0b00000010)
    flags += " SYN";
  if (pkg->tcpFlag & 0b00000100)
    flags += " RST";
  if (pkg->tcpFlag & 0b00001000)
    flags += " PSH";
  if (pkg->tcpFlag & 0b00010000)
    flags += " ACK";
  if (pkg->tcpFlag & 0b00100000)
    flags += " URG";
  log_d("%s: (%d | %s) %s -> %s%s", prefix, pkg->pkgSize, protocol, iptcpaddr(pkg->srcIPv4, pkg->srcPort).c_str(), iptcpaddr(pkg->dstIPv4, pkg->dstPort).c_str(), flags.c_str());
//  c8 buf[65536];
//  memset(buf, '*', 65536);
//  auto size = (ssize_t)jlib::HEX::encode(buf, pkg->in.base(), pkg->in.readable());
//  auto ptr = buf;
//  do {
//    log_v("%c%c %c%c  %c%c %c%c", ptr[0], ptr[1], ptr[2], ptr[3], ptr[4], ptr[5], ptr[6], ptr[7]);
//    size -= 8;
//    ptr += 8;
//  }
//  while (size > 0);
}


ByteBuffer<> Codec::encode (ByteBuffer<>& src, jlib::AES& aes, u8 key[32]) {
  auto hint = jlib::AES::encrypt_size(src.readable());
  debug_assert(hint);
  u8 ivec[16];
  memcpy(ivec, key, 16);
  ByteBuffer<> dst {hint + 5};
  dst.writeByte(PREAMBLE);
  dst.writeByte(PREAMBLE_2);
  auto len = (u16)aes.encrypt(dst.base() + 4, src.readIter(), src.readable(), ivec);
  debug_assert(len > 0);
  dst.writeWord(len);
  dst.move(len);
  dst.writeByte(jlib::SUM::checksum(dst.base() + 2, 2 + len));
  src.skip(src.readable());
  return dst;
}


ByteBuffer<> Codec::decode (ByteBuffer<>& src, jlib::AES& aes, u8 key[32]) {
  while (src.readable()) {
    if (src.readByte() != PREAMBLE) {
      log_d("无效帧起始");
      continue;
    }
    if (src.readable() < 4) {
      src.skip(-1);
      return {};
    }
    if (src.readByte() != PREAMBLE_2) {
      src.skip(-1);
      log_d("无效帧起始 2");
      continue;
    }
    auto len = src.readWord();
    if (len % AES_BLOCK_SIZE) {
      src.skip(-2);
      log_d("载荷未对齐");
      continue;
    }
    if (src.readable() < (uwl)len + 1) {
      src.skip(-4);
      return {};
    }
    auto cs = jlib::SUM::checksum(src.readIter() - 2, len + 2);
    src.skip(len);
    if (src.readByte() != cs) {
      src.skip(-3 - len);
      log_d("帧校验错误");
      continue;
    }
    u8 ivec[16];
    memcpy(ivec, key, 16);
    auto dst = malloc(len);
    aes.decrypt(dst, src.readIter() - 1 - len, len, ivec);
    return {dst, (uwl)len, (uwl)len};
  }
  return {};
}


ByteBuffer<> Codec::encode (ByteBuffer<>& src, jlib::RSA& rsa) {
  auto hint = rsa.public_encrypt_size(src.readable());
  abort_assert(hint);
  ByteBuffer<> dst {hint + 5};
  dst.writeByte(PREAMBLE);
  dst.writeByte(PREAMBLE_2);
  auto len = (u16)rsa.public_encrypt(dst.base() + 4, src.readIter(), src.readable());
  debug_assert(len > 0);
  dst.writeWord(len);
  dst.move(len);
  dst.writeByte(jlib::SUM::checksum(dst.base() + 2, 2 + len));
  return dst;
}


ByteBuffer<> Codec::decode (ByteBuffer<>& src, jlib::RSA& rsa) {
  while (src.readable()) {
    if (src.readByte() != PREAMBLE)
      continue;
    if (src.readable() < 4) {
      src.skip(-1);
      log_e("RSA: 头不足");
      return {};
    }
    if (src.readByte() != PREAMBLE_2) {
      src.skip(-1);
      continue;
    }
    auto len = src.readWord();
    if (src.readable() < (uwl)len + 1) {
      src.skip(-4);
      log_e("RSA: 数据不足");
      return {};
    }
    auto cs = jlib::SUM::checksum(src.readIter() - 2, len + 2);
    src.skip(len);
    if (src.readByte() != cs) {
      src.skip(-3 - len);
      continue;
    }
    auto hint = rsa.private_decrypt_size(len);
    debug_assert(hint);
    auto dst = malloc(hint);
    auto size = (uwl)rsa.private_decrypt(dst, src.readIter() - 1 - len, len);
    return {dst, (uwl)hint, (uwl)size};
  }
  return {};
}












inline static uint16_t loadFragment (const uint8_t*& data, size_t& index, size_t bits) noexcept {
  uint16_t field;
  auto offset = index & 0x7;
  if (offset) {
    auto lowBits = 8U - offset;
    if (lowBits > bits) {
      field = (*data & ((1U << lowBits) - 1U)) >> (lowBits - bits);
      index += bits;
      return field;
    }
    field = *data++ & ((1U << lowBits) - 1U);
    index += lowBits;
    bits -= lowBits;
  }
  else field = 0;
  if (bits >= 8) {
    field = (field << 8U) & *data++;
    index += 8;
    bits -= 8;
  }
  if (bits) {
    field <<= bits;
    field = (field << bits) | (*data >> (8 - bits));
  }
  return field;
}


inline static uint16_t encodeRound (const uint8_t*& data, size_t& index, size_t bits, size_t hinv, size_t roll, size_t linv) noexcept {
  uint16_t field;
  auto offset = index & 0x7;
  if (offset) {
    auto lowBits = 8U - offset;
    if (lowBits > bits) {
      field = (*data & ((1U << lowBits) - 1U)) >> (lowBits - bits);
      index += bits;
      return field;
    }
    field = *data++ & ((1U << lowBits) - 1U);
    index += lowBits;
    bits -= lowBits;
  }
  else field = 0;
  if (bits >= 8) {
    field = (field << 8U) & *data++;
    index += 8;
    bits -= 8;
  }
  if (bits) {
    field <<= bits;
    field = (field << bits) | (*data >> (8 - bits));
  }
  return field;
}